#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <pthread.h>

/*
argv[1] 指定udp，tcp
argv[2] 指定端口号
argv[3] 0关闭心跳 1开启心跳
*/

#define MAX_CLIENTS 5
#define BUFFER_SIZE 1024

typedef struct
{
    int socket_fd;
    struct sockaddr_in address;
} client_info;

int port = 0;

void *handle_client_tcp(void *arg)
{
    client_info client;
    memcpy(&client, arg, sizeof(client_info));
    char buffer[128];
    char buf[INET6_ADDRSTRLEN];
    inet_ntop(AF_INET, &client.address.sin_addr.s_addr, buf, INET6_ADDRSTRLEN);
    printf("%s:%d已通过TCP连接\n", buf, client.address.sin_port);
    while (1)
    {
        memset(&buffer, 0, 128);
        int received_bytes = recv(client.socket_fd, buffer, 128, 0);
        if (received_bytes <= 0)
        {
            // 断开了连接
            close(client.socket_fd);
            printf("%s:%d断开连接\n", buf, client.address.sin_port);
            break;
        }
        printf("收到来自%s:%d的数据%s\n", buf, client.address.sin_port, buffer);
        time_t now;
        now = time(NULL);
        struct tm *local_time;
        local_time = localtime(&now);
        strftime(&buffer[received_bytes], 128 - received_bytes, "(%Y-%m-%d %H:%M:%S server reply)", local_time);

        send(client.socket_fd, buffer, received_bytes + 33 + 1, 0);
    }
    pthread_exit(NULL);
}

void *handle_client_udp(void *arg)
{
    client_info client;
    memcpy(&client, arg, sizeof(client_info));
    char buffer[128];
    char buf[INET6_ADDRSTRLEN];
    socklen_t client_addr_len = sizeof(client.address);
    inet_ntop(AF_INET, &client.address.sin_addr.s_addr, buf, INET6_ADDRSTRLEN);
    printf("%s:%d已通过UDP连接\n", buf, client.address.sin_port);
    while (1)
    {
        memset(&buffer, 0, 128);
        int received_bytes = recvfrom(client.socket_fd, buffer, 128, 0, (struct sockaddr *)&client.address, &client_addr_len);
        if (received_bytes = 0)
        {
            // 断开了连接
            close(client.socket_fd);
            printf("%s:%d断开连接\n", buf, client.address.sin_port);
            break;
        }
        else if (received_bytes < 0)
        {
            perror("Error receiving data");
        }
        printf("收到来自%s:%d的数据%s\n", buf, client.address.sin_port, buffer);
        time_t now;
        now = time(NULL);
        struct tm *local_time;
        local_time = localtime(&now);
        strftime(&buffer[received_bytes], 128 - received_bytes, "(%Y-%m-%d %H:%M:%S server reply)", local_time);

        send(client.socket_fd, buffer, received_bytes + 33 + 1, 0);
    }
    pthread_exit(NULL);
}

int main(int argc, char *argv[])
{
    if (argc < 3 || argc > 4)
    {
        printf("参数数量错误，退出\n");
        return -1;
    }

    // 声明一个socket文件描述符和一个sockaddr_in
    int server_socket_fd, client_socket_fd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    port = atoi(argv[2]);

    if (strcmp(argv[1], "udp") == 0)
    {
        server_socket_fd = socket(AF_INET, SOCK_DGRAM, 0);
        if (server_socket_fd < 0)
        {
            perror("Socket creation failed");
            return 0;
        }

        memset(&server_addr, 0, sizeof(server_addr));
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(port);

        if (bind(server_socket_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)))
        {
            perror("Binding failed");
            return 0;
        }
        printf("使用udp服务,端口号%d\n", port);

        while (1)
        {
            // client_socket_fd = accept(server_socket_fd, (struct sockaddr *)&client_addr, &client_addr_len);
            // if (client_socket_fd < 0)
            // {
            //     perror("Accepting connection failed");
            //     continue;
            // }

            client_info client;
            // client.socket_fd = client_socket_fd;
            // pthread_t tid;
            // pthread_create(&tid, NULL, handle_client_udp, &client);
            char buffer[BUFFER_SIZE];
            memset(&buffer, 0, BUFFER_SIZE);
            int bytes_received = recvfrom(server_socket_fd, buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &client_addr_len);
            memcpy(&(client.address), &client_addr, sizeof(client_addr));

            char buf[INET6_ADDRSTRLEN];
            socklen_t client_addr_len = sizeof(client.address);
            inet_ntop(AF_INET, &client.address.sin_addr.s_addr, buf, INET6_ADDRSTRLEN);

            if (bytes_received == 0)
            {
                printf("%s:%d断开连接\n", buf, client.address.sin_port);
                continue;
            }
            else if (bytes_received < 0)
            {
                continue;
            }
            printf("收到来自%s:%d的数据%s\n", buf, client.address.sin_port, buffer);
            time_t now;
            now = time(NULL);
            struct tm *local_time;
            local_time = localtime(&now);
            strftime(&buffer[bytes_received], BUFFER_SIZE - bytes_received, "(%Y-%m-%d %H:%M:%S server reply)", local_time);
            // inet_pton(AF_INET,"192.168.1.57",&client.address.sin_addr.s_addr);
            sendto(server_socket_fd, buffer, bytes_received + 33 + 1, 0, (struct sockaddr *)&client.address,  client_addr_len);
        }
        close(server_socket_fd);
    }
    else if (strcmp(argv[1], "tcp") == 0)
    {
        server_socket_fd = socket(AF_INET, SOCK_STREAM, 0);
        if (server_socket_fd < 0)
        {
            perror("Socket creation failed");
            return 0;
        }
        memset(&server_addr, 0, sizeof(server_addr));
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(port);

        if (bind(server_socket_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)))
        {
            perror("Binding failed");
            return 0;
        }

        // 监听连接
        printf("使用tcp服务,端口号%d\n", port);
        listen(server_socket_fd, MAX_CLIENTS);
        while (1)
        {
            client_socket_fd = accept(server_socket_fd, (struct sockaddr *)&client_addr, &client_addr_len);
            if (client_socket_fd < 0)
            {
                perror("Accepting connection failed");
                continue;
            }

            client_info client;
            client.socket_fd = client_socket_fd;
            memcpy(&(client.address), &client_addr, sizeof(client_addr));
            pthread_t tid;
            pthread_create(&tid, NULL, handle_client_tcp, &client);
        }
        close(server_socket_fd);
    }
    else
    {
        printf("不支持的参数%s\n", argv[1]);
    }

    return 0;
}